module ub_atomic_alu (
    input logic clk,
    input logic rst_n,
    ldma_ub_dat_if.in u_aru_atomic_rd_dat_if,
    gdma_ub_wr_req_if.in u_aru_atomic_wr_req_if,
    gdma_ub_wr_req_if.out u_aru_wr_req_to_wr_arbiter_if
);

    // Flatten the 2D arrays to 1D for easier processing
    bf16_t rd_dat[`P_ARU * `N0 - 1:0];
    bf16_t wr_dat[`P_ARU * `N0 - 1:0];

    always_comb begin
        for (int i = 0; i < `P_ARU; i++) begin
            for (int j = 0; j < `N0; j++) begin
                rd_dat[i*`N0+j] = u_aru_atomic_rd_dat_if.dat[i][j];
                wr_dat[i*`N0+j] = u_aru_atomic_wr_req_if.pld.dat[i*`N0+j];
            end
        end
    end

    // Both inputs must be valid
    logic both_vld;
    assign both_vld = u_aru_atomic_rd_dat_if.vld && u_aru_atomic_wr_req_if.vld;

    // Backpressure: both ready when output can accept，互锁机制，防止下游rdy，但是队友还没准备好，就接受新数据，相当于这次的rd_dat 被丢弃了
    assign u_aru_atomic_rd_dat_if.rdy = 1'b1;  // RAM数据总是ready，反压通过给fifo反压来实现
    assign u_aru_atomic_wr_req_if.rdy = u_aru_wr_req_to_wr_arbiter_if.rdy && u_aru_atomic_rd_dat_if.vld;    //  RAM的数据vld表示输入的数据肯定已经都准备好了

    // Enable signal for FPU units
    logic comp_enable;
    assign comp_enable = u_aru_wr_req_to_wr_arbiter_if.rdy;

    // Instantiate multipliers for all elements
    bf16_t mul_result   [`P_ARU * `N0 - 1:0];
    logic  mul_out_valid[`P_ARU * `N0 - 1:0];

    genvar i;
    generate
        for (i = 0; i < `P_ARU * `N0; i = i + 1) begin : gen_mul
            bf16_fpmult u_bf16_fpmult (
                .clk         (clk),
                .rst         (~rst_n),
                .clock_enable(comp_enable),
                .in_valid    (both_vld),
                .X           (rd_dat[i]),
                .Y           (wr_dat[i]),
                .R           (mul_result[i]),
                .out_valid   (mul_out_valid[i])
            );
        end
    endgenerate

    // Instantiate adders for all elements
    bf16_t add_result   [`P_ARU * `N0 - 1:0];
    logic  add_out_valid[`P_ARU * `N0 - 1:0];

    generate
        for (i = 0; i < `P_ARU * `N0; i = i + 1) begin : gen_add
            bf16_fpadd u_bf16_fpadd (
                .clk         (clk),
                .rst         (~rst_n),
                .clock_enable(comp_enable),
                .in_valid    (both_vld),
                .X           (rd_dat[i]),
                .Y           (wr_dat[i]),
                .R           (add_result[i]),
                .out_valid   (add_out_valid[i])
            );
        end
    endgenerate

    // Select result based on atomic_mode (delayed to match computation)
    // atomic_mode: 2'b01 - add, 2'b10 - mul
    // Use max latency to align outputs
    localparam MAX_LATENCY = (`ADD_SUB_LATENCY > `MUL_LATENCY) ? `ADD_SUB_LATENCY : `MUL_LATENCY;

    // Delay the atomic_mode and other sideband signals
    gdma_ub_wr_req_s delayed_wr_req;  // 保证控制信号和结果对齐

    common_delay_line #(
        .WIDTH($bits(gdma_ub_wr_req_s)),
        .DEPTH(MAX_LATENCY)
    ) u_sideband_delay (
        .clk     (clk),
        .rst_n   (rst_n),
        .en      (comp_enable),
        .data_in (u_aru_atomic_wr_req_if.pld),
        .data_out(delayed_wr_req)
    );

    // Delay the valid signal
    logic delayed_vld;
    common_delay_line #(
        .WIDTH(1),
        .DEPTH(MAX_LATENCY)
    ) u_vld_delay (
        .clk     (clk),
        .rst_n   (rst_n),
        .en      (comp_enable),  // Always enable for valid signal
        .data_in (both_vld),
        .data_out(delayed_vld)
    );

    // Select between add and mul results based on delayed atomic_mode
    bf16_t selected_result[`P_ARU * `N0 - 1:0];

    always_comb begin
        for (int i = 0; i < `P_ARU * `N0; i = i + 1) begin
            case (delayed_wr_req.atomic_mode)
                2'b01:   selected_result[i] = add_result[i];  // Add
                2'b10:   selected_result[i] = mul_result[i];  // Mul
                default: selected_result[i] = add_result[i];
            endcase
        end
    end

    // Output assignment
    assign u_aru_wr_req_to_wr_arbiter_if.vld = delayed_vld;

    always_comb begin
        u_aru_wr_req_to_wr_arbiter_if.pld = delayed_wr_req;

        // Update lower 64B (P_ARU * N0 elements) with computed results
        for (int i = 0; i < `P_ARU * `N0; i = i + 1) begin
            u_aru_wr_req_to_wr_arbiter_if.pld.dat[i] = selected_result[i];
        end

        // Upper 64B (remaining elements): fill with zeros
        for (int i = `P_ARU * `N0; i < `N0 * `K0; i = i + 1) begin
            u_aru_wr_req_to_wr_arbiter_if.pld.dat[i] = '0;
        end

        // Keep mask unchanged from delayed request (already correct from input)
    end

endmodule
